import torch
import torchvision


class Audio2ParamsNetwork(torch.nn.Module):
    def __init__(self, n_params, n_spaces, pretrained=False):
        super(Audio2ParamsNetwork, self).__init__()
        self.layer = torchvision.models.vgg19_bn(pretrained)

        self.fully_connect1 = torch.nn.Linear(1000, n_params * n_spaces)
        self.fully_connect2 = torch.nn.Linear(n_params * n_spaces, n_params * n_spaces)
        self.fully_connect3 = torch.nn.Linear(n_spaces, n_spaces)

        self.n_params = n_params

    def forward(self, x):
        x = torch.transpose(x, 1, 3)
        x = self.layer(x)
        # x = self.vgg.features(x)
        x = self.fully_connect1(x)
        x = torch.relu(x)
        x = self.fully_connect2(x)
        x = torch.relu(x)
        x = x.view([x.size()[0], self.n_params, -1])
        x = self.fully_connect3(x)
        x = torch.softmax(x, -1)
        return x
